tjbal_unbalance = function(data,  unit, time, outcome, treated, X, estimator, vce){
  
  data = data %>% 
    select(unit, time, treated, outcome, X) 
  colnames(data)[1:4] = c('unit', 'time', 'treated', 'outcome')
  
  baseline_year = data %>% 
    group_by(time) %>% 
    summarise(na_ever = mean(outcome, na.rm = T), 
              is_na_year = as.numeric(is.na(na_ever) == FALSE, 1,0)) %>% 
    filter(is_na_year == 1) %>% 
    summarise(min_year = min(time))
  
  baseline_year = as.numeric(baseline_year)
  
  force_balance = function(data, outcome, baseline_year){
    
    data <- data %>% 
      group_by(unit) %>% 
      filter(time >= baseline_year) %>% 
      mutate(na_ever = mean(outcome)) %>% 
      filter(is.na(na_ever) == F) %>% 
      select(-c(na_ever)) %>% 
      as.data.frame()
    
    return(data)
  }
  
  
  data = force_balance(data = data, outcome = outcome, 
                       baseline_year = baseline_year)
  
  
  
  
  m = tjbal(data = data, Y = 'outcome', D = 'treated', 
            index = c("unit", "time"), 
            X = X,
            demean = F, estimator = estimator, 
            vce = vce)
  return(m)
}


balance_panel_df = readRDS('data/balance_panel_df.rds')

df <- readRDS('data/full_panel_df.rds') %>% 
  filter(cohort == "2011" | cohort == "2012" | cohort == '2013' | cohort == '10000') %>% 
  filter(year >= 2008) %>% 
  filter(year <= 2014) %>% 
  mutate(
    treat = ifelse(treated_ever==1, as.numeric(year>=cohort, 1,0), 0)
  ) %>% 
  select('DISTID', 'year', 'cohort', 'treat', 
         'aog_no_ied', 'IED_Explosion', 'DF', 'ied', 
         'totalcas', 'treated_ever', 'n_fob'
  ) 

anqar_panel = readRDS('data/anqar_panel_2023.rds') %>% 
  filter(year <= 2014)
df = merge(df, anqar_panel, by = c('DISTID', 'year'), all.x = T)

control = read.dta13('data/enumerability/DISTID_cross_section.dta')
df = merge(df,control, by = 'DISTID', all.x = T)
df = df %>% 
  group_by(DISTID) %>% 
  mutate(mean_inf_0810 = mean(gov_influence[year<=2010], na.rm = T ), 
         fob = mean(n_fob[year<=2010]))

tude_list = c('dispute_court', 'return_good', 'gov_influence', 'gov_index2')

att_tjbal <- matrix(, nrow = length(tude_list), ncol = 5)
att_tjbal2 <- matrix(, nrow = length(tude_list), ncol = 5)
att_tjbal3 <- matrix(, nrow = length(tude_list), ncol = 5)

for(j in 1:length(tude_list)){
  
  att_tjbal[j, 1:5] = c(print(tjbal_unbalance(df, 'DISTID', 'year', tude_list[j],
                                            'treat',  'fob', 'mean', 'jack'))[c(1,2,4,5)], tude_list[j])
  
}

for(j in 1:length(tude_list)){
  
  att_tjbal2[j, 1:5] = c(print(tjbal_unbalance(df, 'DISTID', 'year', tude_list[j],
                                              'treat',  'mean_inf_0810', 'mean', 'jack'))[c(1,2,4,5)], tude_list[j])
  
}

df2  = df %>%  
  filter(is.na(not_accessible_taliban) == F)
for(j in 1:length(tude_list)){
  
  att_tjbal3[j, 1:5] =c(print(tjbal_unbalance(df2, 'DISTID', 'year', tude_list[j],
                                             'treat',  'not_accessible_taliban', 'mean', 'jack'))[c(1,2,4,5)], tude_list[j])
  
}
att_tjbal3 = as.data.frame(att_tjbal3) %>% 
  mutate(control_var = as.character('Accessibility'))
att_tjbal2 = as.data.frame(att_tjbal2) %>% 
  mutate(control_var = 'Influence Survey')
att_tjbal = as.data.frame(att_tjbal) %>% 
  mutate(control_var = 'Bases')

plot_df = rbind.data.frame(att_tjbal, att_tjbal2, att_tjbal3)

robust_adjust_control_opinion = ggplot(plot_df, aes(V1, control_var)) + 
  geom_point(size = 2) + 
  geom_errorbar(aes(xmax = V3, xmin = V2), width=.1,
                position=position_dodge(.9)) +
  facet_wrap(~V4, scales = "free", 
             labeller = labeller(V4 = 
                                   c("dispute_court" = "Use Gov. Court",
                                     "gov_index2" = "Gov. Index",
                                     "gov_influence" = "Gov. Influence", 
                                     'return_good' = 'TLB Support'
                                     ))) + 
  xlab('ATT') +
  ylab('Control Variable') + 
  theme_bw()



#### combat #### 

df = df %>% 
  filter(year <= 2013)

out_list = c('aog_no_ied', 'ied', "DF", 
              "IED_Explosion", 'totalcas')

att_tjbal4 <- matrix(, nrow = length(out_list), ncol = 5)
att_tjbal5 <- matrix(, nrow = length(out_list), ncol = 5)
att_tjbal6 <- matrix(, nrow = length(out_list), ncol = 5)

for(j in 1:length(out_list)){
  
  att_tjbal4[j, 1:5] = c(print(tjbal_unbalance(df, 'DISTID', 'year', out_list[j],
                                              'treat',  'fob', 'mean', 'jack'))[c(1,2, 4,5)], out_list[j])
  
}


df2  = df %>%  
  filter(is.na(mean_inf_0810) == F)

for(j in 1:length(out_list)){
  
  att_tjbal5[j, 1:5] = c(print(tjbal_unbalance(df2, 'DISTID', 'year', out_list[j],
                                               'treat',  'mean_inf_0810', 'mean', 'jack'))[c(1,2, 4,5)], out_list[j])
  
}

df2  = df %>%  
  filter(is.na(not_accessible_taliban) == F)

for(j in 1:length(out_list)){
  
  att_tjbal6[j, 1:5] =c(print(tjbal_unbalance(df2, 'DISTID', 'year', out_list[j],
                                              'treat',  'not_accessible_taliban', 'mean', 'jack'))[c(1,2, 4,5)], out_list[j])
  
}
att_tjbal4 = as.data.frame(att_tjbal4) %>% 
  mutate(control_var = as.character('Accessibility'))
att_tjbal5 = as.data.frame(att_tjbal5) %>% 
  mutate(control_var = 'Influence Survey')
att_tjbal6 = as.data.frame(att_tjbal6) %>% 
  mutate(control_var = 'Bases')

plot_df = rbind.data.frame(att_tjbal4, att_tjbal5, att_tjbal6)
plot_df$V1 = as.numeric(plot_df$V1)
plot_df$V2 = as.numeric(plot_df$V2)


robust_adjust_control_combat = ggplot(plot_df, aes(V1, control_var)) + 
  geom_point(size = 2) + 
  geom_errorbar(aes(xmax = V1+1.96*V2, xmin = V1-1.96*V2), width=.1,
                position=position_dodge(.9)) +
  facet_wrap(~V5, scales = "free", 
             labeller = labeller(V5 = 
                                   c("aog_no_ied" = "AOG",
                                     "ied" = "IED (ANSO)",
                                     "IED_Explosion" = "IED (SIGACT)", 
                                     'DF' = 'Direct Fire', 
                                     'totalcas' = 'Total Casualties'
                                   ))
             ) + 
  xlab('ATT') +
  ylab('Control Variable') + 
  theme_bw()


plot_df2 = rbind.data.frame(att_tjbal, att_tjbal2, att_tjbal3,
  att_tjbal4, att_tjbal5, att_tjbal6)
plot_df2$V1 = as.numeric(plot_df2$V1)
plot_df2$V2 = as.numeric(plot_df2$V2)


adjust_control = ggplot(plot_df2, aes(V1, control_var)) + 
  geom_point(size = 2) + 
  geom_errorbar(aes(xmax = V1+1.96*V2, xmin = V1-1.96*V2), width=.1,
                position=position_dodge(.9)) +
  facet_wrap(~V5, scales = "free", 
             labeller = labeller(V5 = 
                                   c("dispute_court" = "Use Gov. Court",
                                     "gov_index2" = "Gov. Index",
                                     "gov_influence" = "Gov. Influence", 
                                     'return_good' = 'TLB Support',
                                     "aog_no_ied" = "AOG",
                                     "ied" = "IED (ANSO)",
                                     "IED_Explosion" = "IED (SIGACT)", 
                                     'DF' = 'Direct Fire', 
                                     'totalcas' = 'Total Casualties'
                                   ))
  ) + 
  xlab('ATT') +
  ylab('Control Variable') + 
  theme_bw()

ggsave('fig-out/adjust_control.pdf', width = 8)

